package com.wotao.wotao.ai.controller;

import org.springframework.ai.chat.ChatClient;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.StreamingChatClient;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.openai.OpenAiChatClient;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.Map;

@RestController
public class ChatController {

    private final ChatClient chatClient;
    private final OpenAiChatClient openAiChatClient;
    private final StreamingChatClient streamingChatClient;

    public ChatController(ChatClient chatClient, OpenAiChatClient openAiChatClient, StreamingChatClient streamingChatClient) {
        this.chatClient = chatClient;
        this.openAiChatClient = openAiChatClient;
        this.streamingChatClient = streamingChatClient;
    }

    @GetMapping("/ai/call")
    public Map call(@RequestParam(value = "message", defaultValue = "中国安徽省黄山市三日游，帮我列个日程") String message) {
        Prompt prompt = new Prompt(new UserMessage(message));
        return Map.of("generation", chatClient.call(message));
    }

    @GetMapping("/ai/generate")
    public String generate(@RequestParam(value = "message", defaultValue = "你好，你会说中文吗") String message) {
        String message1 = "中国安徽省黄山市三日游，帮我列个日程";
        return openAiChatClient.call(message1);
    }

    @GetMapping("/ai/stream")
    public Flux<ChatResponse> stream(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        PromptTemplate promptTemplate = new PromptTemplate(message);
        Prompt prompt = promptTemplate.create(Map.of("adjective", "java", "topic", "冒泡排序"));
        Flux<ChatResponse> rs = openAiChatClient.generateStream(prompt);
        return rs;
    }

    @GetMapping("/ai/generateStream")
    public Flux<ChatResponse> generateStream(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        Prompt prompt = new Prompt(new UserMessage("Tell me a joke"+message));
        return streamingChatClient.generateStream(prompt);
    }
}